package com.albert.websocketsimple;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

@Component
public class MyWebSocketHandler extends TextWebSocketHandler {
	//在线用户列表
	private static final Map<String, WebSocketSession> users = new HashMap<String, WebSocketSession>();

	/**
	 * 连接已关闭，移除在Map集合中的记录
	 */
	@Override
	public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
		users.remove(getUserId(session));
	}

	@Override
	public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
		if (session.isOpen()) {
			session.close();
		}
		System.out.println("连接出错");
		users.remove(getUserId(session));
	}

	/**
	 * 连接建立成功之后，记录用户的连接标识，便于后面发信息
	 */
	@Override
	public void afterConnectionEstablished(WebSocketSession session) throws Exception {
		System.out.println("成功建立连接");
		String userId = getUserId(session);
		System.out.println(userId);
		if (userId != null) {
			users.put(userId, session);
			session.sendMessage(new TextMessage("成功建立socket连接"));
			System.out.println(userId);
			System.out.println(session);
		}
	}

	/**
	 * 处理收到的websocket信息
	 */
	@Override
	protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {

	}

	/**
	* 发送信息给指定用户
	* @param clientId
	* @param message
	* @return
	*/
	public boolean sendMessageToUser(String userId, String message) {
		if (users.get(userId) == null) {
			return false;
		}

		WebSocketSession session = users.get(userId);
		System.out.println("sendMessage:" + session);

		if (!session.isOpen()) {
			return false;
		}
		try {
			int count = 1;
			TextMessage textMessage = null;
			String newMessage = "";

			// 循环向客户端发送数据
			//while (true) {
				newMessage = message + String.valueOf(count);
				textMessage = new TextMessage(message);
				session.sendMessage(textMessage);
				Thread.sleep(500);
				newMessage = "";
			//}
				return true;
		} catch (IOException e) {
			e.printStackTrace();
			return false;
		} catch (InterruptedException e) {
			e.printStackTrace();
			return false;
		}
	}

	/**
	 * 广播信息
	 * @param message
	 * @return
	 */
	public boolean sendMessageToAllUsers(TextMessage message) {
		boolean allSendSuccess = true;
		Set<String> clientIds = users.keySet();
		WebSocketSession session = null;
		for (String clientId : clientIds) {
			try {
				session = users.get(clientId);
				if (session.isOpen()) {
					session.sendMessage(message);
				}
			} catch (IOException e) {
				e.printStackTrace();
				allSendSuccess = false;
			}
		}
		return allSendSuccess;
	}

	/**
	* 获取用户标识
	* @param session
	* @return
	*/
	private String getUserId(WebSocketSession session) {
		try {
			String userId = (String) session.getAttributes().get("user");
			return userId;
		} catch (Exception e) {
			return null;
		}
	}
}
